#include "kernel_operator.h"
using namespace AscendC;
#include <type_traits>
#include <typeinfo>
#include <math.h>
#include <stdbool.h>
#include <vector>
#include <utility>


constexpr int32_t BUFFER_NUM = 2;      //昇腾双buffer技术
 // 定义一个常数作为数组的最大大小
constexpr int32_t  MAX_SPATIAL_DIM = 512;  // 你可以根据需要调整这个大小

class KernelNonMaxSuppression {
public:
    __aicore__ inline KernelNonMaxSuppression() {}

    // Function to compute IoU (Intersection over Union) for two boxes
    __aicore__ inline float compute_iou(float x1, float y1, float x2, float y2, 
                                        float x1_p, float y1_p, float x2_p, float y2_p, 
                                        int center_point_box) {
        float box1_x1, box1_y1, box1_x2, box1_y2;
        float box2_x1, box2_y1, box2_x2, box2_y2;

        //printf("before %f,%f,%f,%f,%f,%f,%f,%f,%d\n ", x1,  y1,  x2,  y2, x1_p,  y1_p,  x2_p,  y2_p, center_point_box);
        if (center_point_box == 0) { // [y1, x1, y2, x2] format
            box1_y1 = y1; box1_x1 = x1; box1_y2 = y2; box1_x2 = x2;
            box2_y1 = y1_p; box2_x1 = x1_p; box2_y2 = y2_p; box2_x2 = x2_p;
        } else { // [x_center, y_center, width, height] format
            box1_x1 = x1 - x2 / 2.0f;
            box1_y1 = y1 - y2 / 2.0f;
            box1_x2 = x1 + x2 / 2.0f;
            box1_y2 = y1 + y2 / 2.0f;

            box2_x1 = x1_p - x2_p / 2.0f;
            box2_y1 = y1_p - y2_p / 2.0f;
            box2_x2 = x1_p + x2_p / 2.0f;
            box2_y2 = y1_p + y2_p / 2.0f;
        }
        //printf("after %f,%f,%f,%f,%f,%f,%f,%f,%d\n ", box1_x1,  box1_y1,  box1_x2,  box1_y2, box2_x1,  box2_y1,  box2_x2,  box2_y2, center_point_box);

        // Replace fmaxf and fminf with conditional operators
        float inter_x1 = (box1_x1 > box2_x1) ? box1_x1 : box2_x1;
        float inter_y1 = (box1_y1 > box2_y1) ? box1_y1 : box2_y1;
        float inter_x2 = (box1_x2 < box2_x2) ? box1_x2 : box2_x2;
        float inter_y2 = (box1_y2 < box2_y2) ? box1_y2 : box2_y2;

        // Calculate the intersection area
        float inter_width = (inter_x2 - inter_x1 > 0) ? (inter_x2 - inter_x1) : 0.0f;
        float inter_height = (inter_y2 - inter_y1 > 0) ? (inter_y2 - inter_y1) : 0.0f;
        float inter_area = inter_width * inter_height;

        // Calculate the area of each box
        float box1_area = (box1_x2 - box1_x1) * (box1_y2 - box1_y1);
        float box2_area = (box2_x2 - box2_x1) * (box2_y2 - box2_y1);
        //printf("inter_area is %f, box1_area is %f, box2_area is %f\n",inter_area, box1_area, box2_area);

        // Compute IoU
        return inter_area / (box1_area + box2_area - inter_area);
    }



    /*
    boxes：包围框，张量（浮点）。形状为  [num_batches、spatial_dimension、4] 的输入张量。单框数据格式由 center_point_box 表示。
    scores：分数，张量（浮点）。形状为 [num_batches、num_classes、spatial_dimension] 的输入张量
    max_output_boxes_per_class（可选）：张量（int64）。表示每个类别每个批次要选择的最大框数的整数。它是一个标量。默认为 0，表示无输出。
    iou_threshold（可选）：张量（浮点数）。浮点数表示用于决定框相对于 IOU 是否重叠过多的阈值。它是标量。值范围 [0, 1]。默认为 0
    score_threshold（可选）：张量（浮点数）。浮点数表示用于根据分数决定何时删除框的阈值。它是一个标量。
    属性：

    center_point_box：int（默认为 0）。整数表示框数据的格式。默认值为 0。0 — 框数据以  [y1, x1, y2, x2] 的形式提供，其中 (y1, x1) 和 (y2, x2) 是任何对角线框角对的坐标，坐标可以作为归一化（即位于区间 [0, 1] 内）或绝对值提供。主要用于 TF 模型。1 — 框数据以 [x_center, y_center, width, height] 的形式提供。主要用于 Pytorch 模型。
    输出：

    selected_indices：张量（int64）。从框张量中选择的索引
    [num_selected_indices，3]，选定的索引格式为 [batch_index，class_index，box_index]。*/

    __aicore__ inline void Init(GM_ADDR boxes,GM_ADDR scores,GM_ADDR max_output_boxes_per_class,GM_ADDR iou_threshold,GM_ADDR score_threshold,GM_ADDR selected_indices, int32_t num_batches, int32_t num_classes, int32_t spatial_dim, int center_point_box)
    {
        this->num_batches = num_batches;
        this->num_classes = num_classes; 
        this->spatial_dim = spatial_dim;
        this->center_point_box = center_point_box; 
        //printf("this->num_batches is %d,this->num_classes is %d, this->spatial_dim is %d, this->center_point_box is %d\n",this->num_batches,this->num_classes,this->spatial_dim,this->center_point_box);

        this->boxes_size = this->num_batches * 4 * this->spatial_dim;
        this->scores_size = this->num_batches * this->num_classes * this->spatial_dim;
        
        bGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_BOXES *>(boxes), this->boxes_size);
        sGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_SCORES *>(scores), this->scores_size);
        mGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_MAX_OUTPUT_BOXES_PER_CLASS *>(max_output_boxes_per_class), 1);
        itGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_IOU_THRESHOLD *>(iou_threshold), 1);
        stGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_SCORE_THRESHOLD *>(score_threshold), 1);
        
        this->mobpc = mGm.GetValue(0);
        this->it = itGm.GetValue(0);
        this->st = stGm.GetValue(0);
        this->total_size = this->mobpc * this->num_batches * this->num_classes;
        siGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_SELECTED_INDICES *>(selected_indices), this->total_size*3);
        //printf("this->mobpc is %d,this->it is %f,this->st is %f, this->total_size is %d\n", this->mobpc,this->it,this->st,this->total_size);


    }
    __aicore__ inline void Process()
    {

    // void non_max_suppression(const float* boxes, const float* scores, int num_batches, int num_classes, int spatial_dim,
    //                       int max_output_boxes_per_class, float iou_threshold, float score_threshold, int center_point_box,
    //                       int64_t* selected_indices) {
    // 初始化 selected_indices，填充为 -1

        for (int i = 0; i < this->total_size; ++i) {
            siGm.SetValue(i * 3, -1);
            siGm.SetValue(i * 3 + 1, -1);
            siGm.SetValue(i * 3 + 2, -1);
        }

    //     // 遍历每个批次和类别
    //     for (int batch = 0; batch < this->num_batches; ++batch) {
    //         for (int cls = 0; cls < this->num_classes; ++cls) {
    //             // 遍历候选框
    //             for (int i = 0; i < this->spatial_dim; ++i) {
    //                 if (sGm.GetValue(batch * this->num_classes * this->spatial_dim + cls * this->spatial_dim + i) < this->st)
    //                     continue;

    //                 // 检查是否与已选框重叠
    //                 bool keep = true;
    //                 int offset = (batch * this->num_classes + cls) * this->mobpc;
    //                 for (int k = 0; k < this->mobpc; ++k) {
    //                     if (siGm.GetValue((offset + k) * 3) == -1)
    //                         break; // 没有更多选框
    //                     int prev_idx = siGm.GetValue((offset + k) * 3 + 2);
    //                     float iou = compute_iou(
    // bGm.GetValue(batch * this->spatial_dim * 4 + i * 4 + 1),  // box1_x1 or box1_x_center
    // bGm.GetValue(batch * this->spatial_dim * 4 + i * 4),      // box1_y1 or box1_y_center
    // bGm.GetValue(batch * this->spatial_dim * 4 + i * 4 + 3),  // box1_x2 or box1_width
    // bGm.GetValue(batch * this->spatial_dim * 4 + i * 4 + 2),  // box1_y2 or box1_height
    // bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4 + 1),  // box2_x1 or box2_x_center
    // bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4),      // box2_y1 or box2_y_center
    // bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4 + 3),  // box2_x2 or box2_width
    // bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4 + 2),  // box2_y2 or box2_height
    // center_point_box);
    //                     printf("batch is %d, cls is %d, i is %d, k is %d, iou is %f\n",batch, cls, i, k, iou);
    //                     if (iou > this->it) {
    //                         keep = false; 
    //                         break;
    //                     }
    //                 }

    //                 // 添加到 selected_indices
    //                 if (keep) {
    //                     for (int k = 0; k < this->mobpc; ++k) {
    //                         if (siGm.GetValue((offset + k) * 3) == -1) {
    //                             siGm.SetValue((offset + k) * 3, batch);
    //                             siGm.SetValue((offset + k) * 3 + 1, cls);
    //                             siGm.SetValue((offset + k) * 3 + 2, i);
    //                             printf("(offset + k) * 3) is %d",(offset + k) * 3);
    //                             break;
    //                         }
    //                     }
    //                 }
    //             }
    //         }
    //     }




        // 创建一个二维数组保存分数和索引
        float score_idx[MAX_SPATIAL_DIM][2] = {0}; // [0] 保存分数，[1] 保存索引
        

        // 遍历每个批次和类别
        for (int batch = 0; batch < this->num_batches; ++batch) {
            for (int cls = 0; cls < this->num_classes; ++cls) {

                // 初始化 score_idx 数组
                for (int i = 0; i < this->spatial_dim; ++i) {
                    score_idx[i][0] = sGm.GetValue(batch * this->num_classes * this->spatial_dim + cls * this->spatial_dim + i);
                    score_idx[i][1] = i; // 保存原始索引
                }

                // 对 score_idx 按分数从大到小排序 (使用简单的冒泡排序实现)
                for (int i = 0; i < this->spatial_dim - 1; ++i) {
                    for (int j = 0; j < this->spatial_dim - i - 1; ++j) {
                        if (score_idx[j][0] < score_idx[j + 1][0]) {
                            // 交换分数和索引
                            float temp_score = score_idx[j][0];
                            float temp_idx = score_idx[j][1];
                            score_idx[j][0] = score_idx[j + 1][0];
                            score_idx[j][1] = score_idx[j + 1][1];
                            score_idx[j + 1][0] = temp_score;
                            score_idx[j + 1][1] = temp_idx;
                        }
                    }
                }

                // 按照排序后的分数顺序处理候选框
                for (int i = 0; i < this->spatial_dim; ++i) {
                    int original_idx = (int)score_idx[i][1];

                    // 获取边界框的坐标
                    float x1 = bGm.GetValue(batch * this->spatial_dim * 4 + original_idx * 4);
                    float y1 = bGm.GetValue(batch * this->spatial_dim * 4 + original_idx * 4 + 1);
                    float x2 = bGm.GetValue(batch * this->spatial_dim * 4 + original_idx * 4 + 2);
                    float y2 = bGm.GetValue(batch * this->spatial_dim * 4 + original_idx * 4 + 3);

                    // 检查 x1, y1, x2, y2 是否有效
                    if (x1 > x2 ) {
                        // 如果无效，跳过当前框
                        //continue;
                        float tmp = x1;
                        x1 = x2;
                        x2 = tmp;
                        bGm.SetValue(batch * this->spatial_dim * 4 + original_idx * 4,   x1);
                        bGm.SetValue(batch * this->spatial_dim * 4 + original_idx * 4+2, x2);
                    }
                    if (y1 > y2 ) {
                        // 如果无效，跳过当前框
                        //continue;
                        float tmp = y1;
                        y1 = y2;
                        y2 = tmp;
                        bGm.SetValue(batch * this->spatial_dim * 4 + original_idx * 4+1, y1);
                        bGm.SetValue(batch * this->spatial_dim * 4 + original_idx * 4+3, y2);
                    }
                    if (x1 == x2 || y1 == y2) continue; 

                    if (score_idx[i][0] < this->st) // 检查分数是否小于阈值
                        continue;

                    // 检查是否与已选框重叠
                    bool keep = true;
                    int offset = (batch * this->num_classes + cls) * this->mobpc;
                    for (int k = 0; k < this->mobpc; ++k) {
                        if (siGm.GetValue((offset + k) * 3) == -1)
                            break; // 没有更多选框

                        int prev_idx = siGm.GetValue((offset + k) * 3 + 2);
                        float iou = compute_iou(
                            y1,x1,y2,x2,
                            bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4 + 1),
                            bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4),
                            bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4 + 3),
                            bGm.GetValue(batch * this->spatial_dim * 4 + prev_idx * 4 + 2),
                            center_point_box);

                        if (iou > this->it) {
                            keep = false;
                            break;
                        }
                    }

                    // 添加到 selected_indices
                    if (keep) {
                        for (int k = 0; k < this->mobpc; ++k) {
                            if (siGm.GetValue((offset + k) * 3) == -1) {
                                siGm.SetValue((offset + k) * 3, batch);
                                siGm.SetValue((offset + k) * 3 + 1, cls);
                                siGm.SetValue((offset + k) * 3 + 2, original_idx);
                                break;
                            }
                        }
                    }
                }
            }
        }

    }



private:
    TPipe pipe;
    //create queue for input, in this case depth is equal to buffer num
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    //create queue for output, in this case depth is equal to buffer num
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;


    GlobalTensor<DTYPE_BOXES> bGm;
    GlobalTensor<DTYPE_SCORES> sGm;
    GlobalTensor<DTYPE_MAX_OUTPUT_BOXES_PER_CLASS> mGm;
    GlobalTensor<DTYPE_IOU_THRESHOLD> itGm;
    GlobalTensor<DTYPE_SCORE_THRESHOLD> stGm;
    GlobalTensor<DTYPE_SELECTED_INDICES> siGm;

    //考生补充自定义成员变量
    int32_t num_batches;
    int32_t num_classes;
    int32_t spatial_dim;
    int32_t center_point_box;

    int32_t boxes_size;
    int32_t scores_size;
    int32_t total_size;
    int32_t mobpc;
    float it;
    float st;



};

extern "C" __global__ __aicore__ void non_max_suppression(GM_ADDR boxes, GM_ADDR scores, GM_ADDR max_output_boxes_per_class, GM_ADDR iou_threshold,GM_ADDR score_threshold,GM_ADDR selected_indices, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelNonMaxSuppression op;  //,DTYPE_INDICES,DTYPE_UPDATES,DTYPE_VAR
    //补充init和process函数调用内容
    op.Init(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, selected_indices, tiling_data.num_batches, tiling_data.num_classes, tiling_data.spatial_dim, tiling_data.center_point_box);
    op.Process();
}

